import argparse
from tqdm import tqdm
from pprint import pprint
import random
from sklearn.cluster import KMeans
import numpy as np

from data.dataset import build_dataset


def parse_args():
    parser = argparse.ArgumentParser(description='clustering anchor size')
    parser.add_argument('--datasets', type=str, help='dataset names', required=True)
    parser.add_argument('--sample', type=int, help='number of sample image, <= 0 for all image', default=0)
    parser.add_argument('--num-anchors', type=int, help='number of anchors', default=45)
    parser.add_argument('--resize', type=str, help='resize size of width,height format', default='512,512')
    args = parser.parse_args()
    return args


def get_boxes_w_h(dataset_names, resize, sample=0):
    dataset_dicts = build_dataset(*dataset_names)
    if sample > 0:
        dataset_dicts = random.sample(dataset_dicts, sample)

    boxes_w_h = []
    for dataset_dict in tqdm(dataset_dicts, desc='extract boxes (w, h)', leave=False):
        image_w, image_h = dataset_dict['params']['width'], dataset_dict['params']['height']
        resize_w, resize_h = resize
        scale = min(resize_w / image_w, resize_h / image_h)
        for obj in dataset_dict['annotations']:
            box_w = (obj['bbox'][2] - obj['bbox'][0]) * scale
            box_h = (obj['bbox'][3] - obj['bbox'][1]) * scale
            boxes_w_h.append([box_w, box_h])

    print("total boxes:", len(boxes_w_h))

    return boxes_w_h


def cluster_boxes(boxes_w_h, num_anchors):
    kmeans = KMeans(n_clusters=num_anchors)
    kmeans.fit(boxes_w_h)
    anchors = kmeans.cluster_centers_
    anchors = np.array(sorted(anchors, key=lambda box_w_h: box_w_h.prod(), reverse=False))
    return anchors


if __name__ == '__main__':
    args = parse_args()
    dataset_names = args.datasets.split(',')
    resize = [float(i) for i in args.resize.split(',')]
    boxes_w_h = get_boxes_w_h(dataset_names, resize, args.sample)
    anchors = cluster_boxes(boxes_w_h, args.num_anchors)
    print(">>> anchors:")
    pprint(anchors)
